// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>
#include <memory>
#include <vector>

#include <nebula/array/builder_base.h>
#include <nebula/array/data.h>

#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>

namespace nebula {

    class TURBO_EXPORT NullBuilder : public ArrayBuilder {
    public:
        explicit NullBuilder(MemoryPool *pool = default_memory_pool(),
                             int64_t TURBO_ARG_UNUSED(alignment) = kDefaultBufferAlignment)
                : ArrayBuilder(pool) {}

        explicit NullBuilder(const std::shared_ptr<DataType> & TURBO_ARG_UNUSED(type),
                             MemoryPool *pool = default_memory_pool(),
                             int64_t alignment = kDefaultBufferAlignment)
                : NullBuilder(pool, alignment) {}

        /// \brief append the specified number of null elements
        turbo::Status append_nulls(int64_t length) final {
            if (length < 0) return turbo::invalid_argument_error("length must be positive");
            null_count_ += length;
            length_ += length;
            return turbo::OkStatus();
        }

        /// \brief append a single null element
        turbo::Status append_null() final { return append_nulls(1); }

        turbo::Status append_empty_values(int64_t length) final { return append_nulls(length); }

        turbo::Status append_empty_value() final { return append_empty_values(1); }

        turbo::Status append(std::nullptr_t) { return append_null(); }

        turbo::Status append_array_slice(const ArraySpan &, int64_t, int64_t length) override {
            return append_nulls(length);
        }

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override;

        /// \cond FALSE
        using ArrayBuilder::finish;
        /// \endcond

        std::shared_ptr<DataType> type() const override { return null(); }

        turbo::Status finish(std::shared_ptr<NullArray> *out) { return finish_typed(out); }
    };

    /// \addtogroup numeric-builders
    ///
    /// @{

    /// Base class for all Builders that emit an Array of a scalar numerical type.
    template<typename T>
    class NumericBuilder
            : public ArrayBuilder,
              public internal::ArrayBuilderExtraOps<NumericBuilder<T>, typename T::c_type> {
    public:
        using TypeClass = T;
        using value_type = typename T::c_type;
        using ArrayType = typename TypeTraits<T>::ArrayType;

        template<typename T1 = T>
        explicit NumericBuilder(
                enable_if_parameter_free<T1, MemoryPool *> pool = default_memory_pool(),
                int64_t alignment = kDefaultBufferAlignment)
                : ArrayBuilder(pool, alignment),
                  type_(TypeTraits<T>::type_singleton()),
                  data_builder_(pool, alignment) {}

        NumericBuilder(const std::shared_ptr<DataType> &type, MemoryPool *pool,
                       int64_t alignment = kDefaultBufferAlignment)
                : ArrayBuilder(pool, alignment), type_(type), data_builder_(pool, alignment) {}

        /// append a single scalar and increase the size if necessary.
        turbo::Status append(const value_type val) {
            TURBO_RETURN_NOT_OK(ArrayBuilder::Reserve(1));
            unsafe_append(val);
            return turbo::OkStatus();
        }

        /// write nulls as uint8_t* (0 value indicates null) into pre-allocated memory
        /// The memory at the corresponding data slot is set to 0 to prevent
        /// uninitialized memory access
        turbo::Status append_nulls(int64_t length) final {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(length, value_type{});  // zero
            unsafe_set_null(length);
            return turbo::OkStatus();
        }

        /// \brief append a single null element
        turbo::Status append_null() final {
            TURBO_RETURN_NOT_OK(Reserve(1));
            data_builder_.unsafe_append(value_type{});  // zero
            unsafe_append_to_bitmap(false);
            return turbo::OkStatus();
        }

        /// \brief append a empty element
        turbo::Status append_empty_value() final {
            TURBO_RETURN_NOT_OK(Reserve(1));
            data_builder_.unsafe_append(value_type{});  // zero
            unsafe_append_to_bitmap(true);
            return turbo::OkStatus();
        }

        /// \brief append several empty elements
        turbo::Status append_empty_values(int64_t length) final {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(length, value_type{});  // zero
            unsafe_set_not_null(length);
            return turbo::OkStatus();
        }

        value_type get_value(int64_t index) const { return data_builder_.data()[index]; }

        void reset() override {
            data_builder_.reset();
            ArrayBuilder::reset();
        }

        turbo::Status resize(int64_t capacity) override {
            TURBO_RETURN_NOT_OK(check_capacity(capacity));
            capacity = std::max(capacity, kMinBuilderCapacity);
            TURBO_RETURN_NOT_OK(data_builder_.resize(capacity));
            return ArrayBuilder::resize(capacity);
        }

        value_type operator[](int64_t index) const { return get_value(index); }

        value_type &operator[](int64_t index) {
            return reinterpret_cast<value_type *>(data_builder_.mutable_data())[index];
        }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous C array of values
        /// \param[in] length the number of values to append
        /// \param[in] valid_bytes an optional sequence of bytes where non-zero
        /// indicates a valid (non-null) value
        /// \return turbo::Status
        turbo::Status append_values(const value_type *values, int64_t length,
                                   const uint8_t *valid_bytes = nullptr) {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(values, length);
            // length_ is update by these
            ArrayBuilder::unsafe_append_to_bitmap(valid_bytes, length);
            return turbo::OkStatus();
        }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous C array of values
        /// \param[in] length the number of values to append
        /// \param[in] bitmap a validity bitmap to copy (may be null)
        /// \param[in] bitmap_offset an offset into the validity bitmap
        /// \return turbo::Status
        turbo::Status append_values(const value_type *values, int64_t length, const uint8_t *bitmap,
                                   int64_t bitmap_offset) {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(values, length);
            // length_ is update by these
            ArrayBuilder::unsafe_append_to_bitmap(bitmap, bitmap_offset, length);
            return turbo::OkStatus();
        }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous C array of values
        /// \param[in] length the number of values to append
        /// \param[in] is_valid an std::vector<bool> indicating valid (1) or null
        /// (0). Equal in length to values
        /// \return turbo::Status
        turbo::Status append_values(const value_type *values, int64_t length,
                                   const std::vector<bool> &is_valid) {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(values, length);
            // length_ is update by these
            ArrayBuilder::unsafe_append_to_bitmap(is_valid);
            return turbo::OkStatus();
        }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a std::vector of values
        /// \param[in] is_valid an std::vector<bool> indicating valid (1) or null
        /// (0). Equal in length to values
        /// \return turbo::Status
        turbo::Status append_values(const std::vector<value_type> &values,
                                   const std::vector<bool> &is_valid) {
            return append_values(values.data(), static_cast<int64_t>(values.size()), is_valid);
        }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a std::vector of values
        /// \return turbo::Status
        turbo::Status append_values(const std::vector<value_type> &values) {
            return append_values(values.data(), static_cast<int64_t>(values.size()));
        }

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override {
            TURBO_MOVE_OR_RAISE(auto null_bitmap,
                                null_bitmap_builder_.FinishWithLength(length_));
            TURBO_MOVE_OR_RAISE(auto data, data_builder_.FinishWithLength(length_));
            *out = ArrayData::create(type(), length_, {null_bitmap, data}, null_count_);
            capacity_ = length_ = null_count_ = 0;
            return turbo::OkStatus();
        }

        /// \cond FALSE
        using ArrayBuilder::finish;
        /// \endcond

        turbo::Status finish(std::shared_ptr<ArrayType> *out) { return finish_typed(out); }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values_begin InputIterator to the beginning of the values
        /// \param[in] values_end InputIterator pointing to the end of the values
        /// \return turbo::Status
        template<typename ValuesIter>
        turbo::Status append_values(ValuesIter values_begin, ValuesIter values_end) {
            int64_t length = static_cast<int64_t>(std::distance(values_begin, values_end));
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(values_begin, values_end);
            // this updates the length_
            unsafe_set_not_null(length);
            return turbo::OkStatus();
        }

        /// \brief append a sequence of elements in one shot, with a specified nullmap
        /// \param[in] values_begin InputIterator to the beginning of the values
        /// \param[in] values_end InputIterator pointing to the end of the values
        /// \param[in] valid_begin InputIterator with elements indication valid(1)
        ///  or null(0) values.
        /// \return turbo::Status
        template<typename ValuesIter, typename ValidIter>
        enable_if_t<!std::is_pointer<ValidIter>::value, turbo::Status> append_values(
                ValuesIter values_begin, ValuesIter values_end, ValidIter valid_begin) {
            static_assert(!internal::is_null_pointer<ValidIter>::value,
                          "Don't pass a nullptr directly as valid_begin, use the 2-argument "
                          "version instead");
            int64_t length = static_cast<int64_t>(std::distance(values_begin, values_end));
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(values_begin, values_end);
            null_bitmap_builder_.unsafe_append<true>(
                    length, [&valid_begin]() -> bool { return *valid_begin++; });
            length_ = null_bitmap_builder_.length();
            null_count_ = null_bitmap_builder_.false_count();
            return turbo::OkStatus();
        }

        // Same as above, with a pointer type ValidIter
        template<typename ValuesIter, typename ValidIter>
        enable_if_t<std::is_pointer<ValidIter>::value, turbo::Status> append_values(
                ValuesIter values_begin, ValuesIter values_end, ValidIter valid_begin) {
            int64_t length = static_cast<int64_t>(std::distance(values_begin, values_end));
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(values_begin, values_end);
            // this updates the length_
            if (valid_begin == nullptr) {
                unsafe_set_not_null(length);
            } else {
                null_bitmap_builder_.unsafe_append<true>(
                        length, [&valid_begin]() -> bool { return *valid_begin++; });
                length_ = null_bitmap_builder_.length();
                null_count_ = null_bitmap_builder_.false_count();
            }

            return turbo::OkStatus();
        }

        turbo::Status append_array_slice(const ArraySpan &array, int64_t offset,
                                       int64_t length) override {
            return append_values(array.get_values<value_type>(1) + offset, length,
                                array.get_values<uint8_t>(0, 0), array.offset + offset);
        }

        /// append a single scalar under the assumption that the underlying Buffer is
        /// large enough.
        ///
        /// This method does not capacity-check; make sure to call Reserve
        /// beforehand.
        void unsafe_append(const value_type val) {
            ArrayBuilder::unsafe_append_to_bitmap(true);
            data_builder_.unsafe_append(val);
        }

        void unsafe_append_null() {
            ArrayBuilder::unsafe_append_to_bitmap(false);
            data_builder_.unsafe_append(value_type{});  // zero
        }

        std::shared_ptr<DataType> type() const override { return type_; }

    protected:
        std::shared_ptr<DataType> type_;
        TypedBufferBuilder<value_type> data_builder_;
    };

    // Builders

    using UInt8Builder = NumericBuilder<UInt8Type>;
    using UInt16Builder = NumericBuilder<UInt16Type>;
    using UInt32Builder = NumericBuilder<UInt32Type>;
    using UInt64Builder = NumericBuilder<UInt64Type>;

    using Int8Builder = NumericBuilder<Int8Type>;
    using Int16Builder = NumericBuilder<Int16Type>;
    using Int32Builder = NumericBuilder<Int32Type>;
    using Int64Builder = NumericBuilder<Int64Type>;

    using Fp16Builder = NumericBuilder<Fp16Type>;
    using Fp32Builder = NumericBuilder<Fp32Type>;
    using Fp64Builder = NumericBuilder<Fp64Type>;

    /// @}

    /// \addtogroup temporal-builders
    ///
    /// @{

    using Date32Builder = NumericBuilder<Date32Type>;
    using Date64Builder = NumericBuilder<Date64Type>;
    using Time32Builder = NumericBuilder<Time32Type>;
    using Time64Builder = NumericBuilder<Time64Type>;
    using TimestampBuilder = NumericBuilder<TimestampType>;
    using MonthIntervalBuilder = NumericBuilder<MonthIntervalType>;
    using DurationBuilder = NumericBuilder<DurationType>;

    /// @}

    class TURBO_EXPORT BooleanBuilder
            : public ArrayBuilder,
              public internal::ArrayBuilderExtraOps<BooleanBuilder, bool> {
    public:
        using TypeClass = BooleanType;
        using value_type = bool;

        explicit BooleanBuilder(MemoryPool *pool = default_memory_pool(),
                                int64_t alignment = kDefaultBufferAlignment);

        BooleanBuilder(const std::shared_ptr<DataType> &type,
                       MemoryPool *pool = default_memory_pool(),
                       int64_t alignment = kDefaultBufferAlignment);

        /// write nulls as uint8_t* (0 value indicates null) into pre-allocated memory
        turbo::Status append_nulls(int64_t length) final {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(length, false);
            unsafe_set_null(length);
            return turbo::OkStatus();
        }

        turbo::Status append_null() final {
            TURBO_RETURN_NOT_OK(Reserve(1));
            unsafe_append_null();
            return turbo::OkStatus();
        }

        turbo::Status append_empty_value() final {
            TURBO_RETURN_NOT_OK(Reserve(1));
            data_builder_.unsafe_append(false);
            unsafe_set_not_null(1);
            return turbo::OkStatus();
        }

        turbo::Status append_empty_values(int64_t length) final {
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append(length, false);
            unsafe_set_not_null(length);
            return turbo::OkStatus();
        }

        /// Scalar append
        turbo::Status append(const bool val) {
            TURBO_RETURN_NOT_OK(Reserve(1));
            unsafe_append(val);
            return turbo::OkStatus();
        }

        turbo::Status append(const uint8_t val) { return append(val != 0); }

        /// Scalar append, without checking for capacity
        void unsafe_append(const bool val) {
            data_builder_.unsafe_append(val);
            unsafe_append_to_bitmap(true);
        }

        void unsafe_append_null() {
            data_builder_.unsafe_append(false);
            unsafe_append_to_bitmap(false);
        }

        void unsafe_append(const uint8_t val) { unsafe_append(val != 0); }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous array of bytes (non-zero is 1)
        /// \param[in] length the number of values to append
        /// \param[in] valid_bytes an optional sequence of bytes where non-zero
        /// indicates a valid (non-null) value
        /// \return turbo::Status
        turbo::Status append_values(const uint8_t *values, int64_t length,
                                   const uint8_t *valid_bytes = nullptr);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a bitmap of values
        /// \param[in] length the number of values to append
        /// \param[in] validity a validity bitmap to copy (may be null)
        /// \param[in] offset an offset into the values and validity bitmaps
        /// \return turbo::Status
        turbo::Status append_values(const uint8_t *values, int64_t length, const uint8_t *validity,
                                   int64_t offset);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous C array of values
        /// \param[in] length the number of values to append
        /// \param[in] is_valid an std::vector<bool> indicating valid (1) or null
        /// (0). Equal in length to values
        /// \return turbo::Status
        turbo::Status append_values(const uint8_t *values, int64_t length,
                                   const std::vector<bool> &is_valid);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a std::vector of bytes
        /// \param[in] is_valid an std::vector<bool> indicating valid (1) or null
        /// (0). Equal in length to values
        /// \return turbo::Status
        turbo::Status append_values(const std::vector<uint8_t> &values,
                                   const std::vector<bool> &is_valid);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a std::vector of bytes
        /// \return turbo::Status
        turbo::Status append_values(const std::vector<uint8_t> &values);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values an std::vector<bool> indicating true (1) or false
        /// \param[in] is_valid an std::vector<bool> indicating valid (1) or null
        /// (0). Equal in length to values
        /// \return turbo::Status
        turbo::Status append_values(const std::vector<bool> &values, const std::vector<bool> &is_valid);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values an std::vector<bool> indicating true (1) or false
        /// \return turbo::Status
        turbo::Status append_values(const std::vector<bool> &values);

        /// \brief append a sequence of elements in one shot
        /// \param[in] values_begin InputIterator to the beginning of the values
        /// \param[in] values_end InputIterator pointing to the end of the values
        ///  or null(0) values
        /// \return turbo::Status
        template<typename ValuesIter>
        turbo::Status append_values(ValuesIter values_begin, ValuesIter values_end) {
            int64_t length = static_cast<int64_t>(std::distance(values_begin, values_end));
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append<false>(
                    length, [&values_begin]() -> bool { return *values_begin++; });
            // this updates length_
            unsafe_set_not_null(length);
            return turbo::OkStatus();
        }

        /// \brief append a sequence of elements in one shot, with a specified nullmap
        /// \param[in] values_begin InputIterator to the beginning of the values
        /// \param[in] values_end InputIterator pointing to the end of the values
        /// \param[in] valid_begin InputIterator with elements indication valid(1)
        ///  or null(0) values
        /// \return turbo::Status
        template<typename ValuesIter, typename ValidIter>
        enable_if_t<!std::is_pointer<ValidIter>::value, turbo::Status> append_values(
                ValuesIter values_begin, ValuesIter values_end, ValidIter valid_begin) {
            static_assert(!internal::is_null_pointer<ValidIter>::value,
                          "Don't pass a nullptr directly as valid_begin, use the 2-argument "
                          "version instead");
            int64_t length = static_cast<int64_t>(std::distance(values_begin, values_end));
            TURBO_RETURN_NOT_OK(Reserve(length));

            data_builder_.unsafe_append<false>(
                    length, [&values_begin]() -> bool { return *values_begin++; });
            null_bitmap_builder_.unsafe_append<true>(
                    length, [&valid_begin]() -> bool { return *valid_begin++; });
            length_ = null_bitmap_builder_.length();
            null_count_ = null_bitmap_builder_.false_count();
            return turbo::OkStatus();
        }

        // Same as above, for a pointer type ValidIter
        template<typename ValuesIter, typename ValidIter>
        enable_if_t<std::is_pointer<ValidIter>::value, turbo::Status> append_values(
                ValuesIter values_begin, ValuesIter values_end, ValidIter valid_begin) {
            int64_t length = static_cast<int64_t>(std::distance(values_begin, values_end));
            TURBO_RETURN_NOT_OK(Reserve(length));
            data_builder_.unsafe_append<false>(
                    length, [&values_begin]() -> bool { return *values_begin++; });

            if (valid_begin == nullptr) {
                unsafe_set_not_null(length);
            } else {
                null_bitmap_builder_.unsafe_append<true>(
                        length, [&valid_begin]() -> bool { return *valid_begin++; });
            }
            length_ = null_bitmap_builder_.length();
            null_count_ = null_bitmap_builder_.false_count();
            return turbo::OkStatus();
        }

        turbo::Status append_values(int64_t length, bool value);

        turbo::Status append_array_slice(const ArraySpan &array, int64_t offset,
                                       int64_t length) override {
            return append_values(array.get_values<uint8_t>(1, 0), length,
                                array.get_values<uint8_t>(0, 0), array.offset + offset);
        }

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override;

        /// \cond FALSE
        using ArrayBuilder::finish;
        /// \endcond

        turbo::Status finish(std::shared_ptr<BooleanArray> *out) { return finish_typed(out); }

        void reset() override;

        turbo::Status resize(int64_t capacity) override;

        std::shared_ptr<DataType> type() const override { return boolean(); }

    protected:
        TypedBufferBuilder<bool> data_builder_;
    };

}  // namespace nebula
